Introduction

This notebook is used to generate the various discrete PMFs in the lecture notes.


In [1]:
from __future__ import division

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl

mpl.rcParams['lines.linewidth'] = 2
mpl.rcParams['lines.color'] = 'r'
mpl.rcParams['axes.titlesize'] = 32
mpl.rcParams['axes.labelsize'] = 24 
mpl.rcParams['axes.labelsize'] = 24 
mpl.rcParams['xtick.labelsize'] = 24 
mpl.rcParams['ytick.labelsize'] = 24 

%matplotlib inline

In [2]:
def make_bar_plot(x, y, xlabel="$x$", ylabel="$p_X(x)$"):
    """
    Utility function to generate bar plots from x and y arrays.
    """
    plt.figure(figsize=(12, 8))
    plt.bar(x, y, width=0.2, color="k")
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.xlim([min(x)-1, max(x)+1])
    return plt

In [3]:
# Roll of two four-sided dice
plt = make_bar_plot(np.arange(2, 9), [1/16, 1/8, 3/16, 4/16, 3/16, 1/8, 1/16])



In [4]:
#CDF of roll of two four-sided dice
make_bar_plot(np.arange(2, 9), np.cumsum([1/16, 1/8, 3/16, 4/16, 3/16, 1/8, 1/16]), ylabel="$F_X(x)$")


Out[4]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [5]:
#Uniform distribution
make_bar_plot(np.arange(-3, 4), [1/7] * 7)


Out[5]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [6]:
# Z=X**2 of uniform PMF
make_bar_plot(np.arange(0, 10), [1/7, 2/7, 0, 0, 2/7, 0, 0, 0, 0, 2/7], ylabel="$p_Z(z)$", xlabel="$z$")


Out[6]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [7]:
#Binomial
from scipy.stats import binom
n = 10
p = 0.5
x = np.arange(0, n + 1)
make_bar_plot(x, binom.pmf(x, n, p), xlabel="$k", ylabel="$p_X(k)$")


Out[7]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [8]:
n = 10
p = 0.5
x = np.arange(0, n + 1)
make_bar_plot(x, binom.pmf(x, n, p), xlabel="$k$", ylabel="$p_X(k)$")


Out[8]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [9]:
# Geometric
p = 0.5
x = np.arange(1, 10)
make_bar_plot(x, p * (1 - p) ** x, xlabel="$k$", ylabel="$p_X(k)$")


Out[9]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [10]:
from scipy.stats import poisson
x = np.arange(100)
l = 10
make_bar_plot(x, poisson.pmf(x, l), xlabel="$k$", ylabel="$p_X(k)$")


Out[10]:
<module 'matplotlib.pyplot' from '/Users/shyue/miniconda3/lib/python3.6/site-packages/matplotlib/pyplot.py'>

In [14]:
from scipy.stats import nbinom
from scipy.misc import factorial
r = 3
x = np.arange(0, 20)
p = 0.5
coeffs = factorial(x + r - 1) / factorial(x) / factorial(r - 1)

plt = make_bar_plot(x, coeffs * (1-p) ** r * p ** x, xlabel="$k$", ylabel="$p_X(k)$")
plt.xlim([0, 10])


Out[14]:
(0, 10)

In [12]:
x = np.arange(r, 20)
plt = make_bar_plot(x, nbinom.pmf(x, 3, 0.5), xlabel="$k$", ylabel="$p_X(k)$")



In [13]:
from scipy.stats import binom
n = 10
p = 0.9
x = np.arange(0, 100)
plt = make_bar_plot(x, binom.pmf(x, n, p), xlabel="$k$", ylabel="$p_X(k)$")



In [ ]: